package utils;

import com.ansj.vec.Learn;
import com.ansj.vec.Word2VEC;

import java.io.File;
import java.io.IOException;

/**
 * Created by 梅晨 on 2017/7/20.
 */
public class VecModel {
    private String modelPath;
    private String dicPath;

    public VecModel(String modelPath, String dicPath) {
        this.modelPath = modelPath;
        this.dicPath = dicPath;
    }

    /**
     * 根据worddic文件训练word2vec模型
     */
    public void train()throws IOException{
        Learn learn = new Learn();
        learn.learnFile(new File(dicPath));
        learn.saveModel(new File(modelPath));
    }

    /**
     * 从modelPath加载word2vec模型
     * @return
     */
    public Word2VEC loadModel()throws IOException{
        Word2VEC w2v = new Word2VEC();
        w2v.loadJavaModel(modelPath);
        return w2v;
    }
}
